import pymysql
import pymysql.cursors
import datetime
import getpass
from src.core.config import get_database_config

class MySQLDBGenerator:
    def __init__(self):
        # 从配置读取器获取数据库配置
        db_config = get_database_config()
        self.host = db_config.get('host', 'localhost')
        self.port = db_config.get('port', 3306)
        self.user = db_config.get('user') or input('请输入MySQL用户名: ')
        self.password = db_config.get('password') or getpass.getpass('请输入MySQL密码: ')
        self.db_name = db_config.get('db_name', 'asset')
        self.connection = None
        self.cursor = None
        
        # 定义数据库表结构
        self.tables = {
            'owners': {
                'columns': [
                    ('owner_id', 'INT PRIMARY KEY AUTO_INCREMENT'),
                    ('owner_name', 'VARCHAR(100) NOT NULL'),
                    ('create_time', 'DATETIME NOT NULL'),
                    ('update_time', 'DATETIME NOT NULL')
                ]
            },
            'owner_relations': {
                'columns': [
                    ('id', 'INT PRIMARY KEY AUTO_INCREMENT'),
                    ('parent_id', 'INT NOT NULL'),
                    ('child_id', 'INT NOT NULL'),
                    ('parent_type', 'VARCHAR(50)'),
                    ('child_type', 'VARCHAR(50)'),
                    ('create_time', 'DATETIME NOT NULL'),
                    ('update_time', 'DATETIME NOT NULL')
                ]
            },
            'ips': {
                'columns': [
                    ('ip_id', 'INT PRIMARY KEY AUTO_INCREMENT'),
                    ('ip', 'VARCHAR(50) NOT NULL'),
                    ('owner_id', 'INT'),
                    # 是否可以ping通
                    ('is_pingable', 'TINYINT(1) NOT NULL DEFAULT 0'),
                    ('scan_status', 'TINYINT(1) NOT NULL DEFAULT 0'), # 扫描状态，只有所有端口都扫描过才为1
                    # 端口扫描量
                    ('port_scan_count', 'INT NOT NULL DEFAULT 0'),
                    ('last_scan_time', 'DATETIME'), # 最后扫描时间
                    ('create_time', 'DATETIME NOT NULL'),
                    ('update_time', 'DATETIME NOT NULL')
                ]
            },
            'domains': {
                'columns': [
                    ('domain_id', 'INT PRIMARY KEY AUTO_INCREMENT'),
                    ('domain', 'VARCHAR(255) NOT NULL'),
                    ('owner_id', 'INT'),
                    ('create_time', 'DATETIME NOT NULL'),
                    ('update_time', 'DATETIME NOT NULL')
                ]
            },
            'ipdomains': {
                'columns': [
                    ('id', 'INT PRIMARY KEY AUTO_INCREMENT'),
                    ('ip', 'VARCHAR(50) NOT NULL'),
                    ('domain', 'VARCHAR(255) NOT NULL'),
                    ('is_valid', 'TINYINT(1) NOT NULL DEFAULT 0'), # 是否有效
                    ('last_check_time', 'DATETIME'), # 最后检测时间
                    ('create_time', 'DATETIME NOT NULL'), # 创建时间
                    ('update_time', 'DATETIME NOT NULL') # 更新时间
                ]
            },
            'assets': {
                'columns': [
                    ('asset_id', 'INT PRIMARY KEY AUTO_INCREMENT'),
                    ('owner_id', 'INT'),
                    ('owner_name', 'VARCHAR(255)'),
                    ('asset_name', 'VARCHAR(255)'),
                    ('ip', 'VARCHAR(50)'),
                    ('port', 'INT NOT NULL'),
                    ('domain', 'VARCHAR(255)'),
                    ('asset_protocol', 'VARCHAR(50) NOT NULL'),
                    ('asset_service', 'VARCHAR(50)'),
                    ('asset_fingerprint', 'TEXT'),
                    ('asset_title', 'VARCHAR(255)'),
                    ('status_code', 'VARCHAR(50)'),
                    ('asset_response', 'LONGTEXT'),
                    ('asset_tags', 'VARCHAR(255)'),
                    ('is_valid', 'TINYINT(1) NOT NULL DEFAULT 0'), # 是否有效
                    ('last_check_time', 'DATETIME'), # 最后检测时间
                    ('create_time', 'DATETIME NOT NULL'), # 创建时间
                    ('update_time', 'DATETIME NOT NULL'), # 更新时间
                ]
            }
        }
    
    def connect(self, use_db=False):
        """连接到MySQL数据库"""
        try:
            db_name = self.db_name if use_db else None
            self.connection = pymysql.connect(
                host=self.host,
                port=self.port,
                user=self.user,
                password=self.password,
                database=db_name,
                cursorclass=pymysql.cursors.DictCursor
            )
            self.cursor = self.connection.cursor()
            print(f"成功连接到MySQL{'数据库 ' + self.db_name if use_db else '服务器'}")
            return True
        except pymysql.MySQLError as e:
            print(f"连接失败: {e}")
            return False
    
    def disconnect(self):
        """断开与MySQL数据库的连接"""
        if self.connection:
            self.cursor.close()
            self.connection.close()
            print("已断开与MySQL数据库的连接")
    
    def check_database_exists(self):
        """检查数据库是否存在"""
        try:
            self.cursor.execute("SHOW DATABASES LIKE %s", (self.db_name,))
            return self.cursor.fetchone() is not None
        except pymysql.MySQLError as e:
            print(f"检查数据库失败: {e}")
            return False
    
    def create_database(self):
        """创建数据库"""
        try:
            self.cursor.execute(f"CREATE DATABASE IF NOT EXISTS {self.db_name} CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci")
            self.connection.commit()
            print(f"数据库 {self.db_name} 创建成功")
            return True
        except pymysql.MySQLError as e:
            print(f"创建数据库失败: {e}")
            self.connection.rollback()
            return False
    
    def check_table_exists(self, table_name):
        """检查表是否存在"""
        try:
            self.cursor.execute("SHOW TABLES LIKE %s", (table_name,))
            return self.cursor.fetchone() is not None
        except pymysql.MySQLError as e:
            print(f"检查表 {table_name} 失败: {e}")
            return False
    
    def get_table_columns(self, table_name):
        """获取表的列信息"""
        try:
            self.cursor.execute(f"DESCRIBE {table_name}")
            return self.cursor.fetchall()
        except pymysql.MySQLError as e:
            print(f"获取表 {table_name} 列信息失败: {e}")
            return []
    
    def create_table(self, table_name, columns):
        """创建表"""
        try:
            columns_sql = ', '.join([f"{col_name} {col_type}" for col_name, col_type in columns])
            create_sql = f"CREATE TABLE IF NOT EXISTS {table_name} ({columns_sql})"
            self.cursor.execute(create_sql)
            self.connection.commit()
            print(f"表 {table_name} 创建成功")
            return True
        except pymysql.MySQLError as e:
            print(f"创建表 {table_name} 失败: {e}")
            self.connection.rollback()
            return False
    
    def add_column(self, table_name, column_name, column_type):
        """添加列"""
        try:
            alter_sql = f"ALTER TABLE {table_name} ADD COLUMN {column_name} {column_type}"
            self.cursor.execute(alter_sql)
            self.connection.commit()
            print(f"表 {table_name} 添加列 {column_name} 成功")
            return True
        except pymysql.MySQLError as e:
            print(f"表 {table_name} 添加列 {column_name} 失败: {e}")
            self.connection.rollback()
            return False
    
    def modify_column(self, table_name, column_name, column_type):
        """修改列"""
        try:
            alter_sql = f"ALTER TABLE {table_name} MODIFY COLUMN {column_name} {column_type}"
            self.cursor.execute(alter_sql)
            self.connection.commit()
            print(f"表 {table_name} 修改列 {column_name} 成功")
            return True
        except pymysql.MySQLError as e:
            print(f"表 {table_name} 修改列 {column_name} 失败: {e}")
            self.connection.rollback()
            return False
    
    def generate(self):
        """生成数据库和表"""
        # 连接到MySQL服务器
        if not self.connect(use_db=False):
            return False
        
        # 检查并创建数据库
        if not self.check_database_exists():
            if not self.create_database():
                self.disconnect()
                return False
        
        # 断开连接，重新连接到指定数据库
        self.disconnect()
        if not self.connect(use_db=True):
            return False
        
        # 处理每个表
        for table_name, table_info in self.tables.items():
            columns = table_info['columns']
            
            # 检查并创建表
            if not self.check_table_exists(table_name):
                if not self.create_table(table_name, columns):
                    continue
            else:
                # 获取现有表的列信息
                existing_columns = self.get_table_columns(table_name)
                existing_columns_dict = {col['Field']: col['Type'] for col in existing_columns}
                
                # 检查列是否存在，不存在则添加
                for column_def in columns:
                    # 所有列定义已统一为元组格式
                    if isinstance(column_def, tuple):
                        col_name, col_type = column_def
                    else:
                        continue

                    # 获取完整的预期类型（不包含额外约束如NOT NULL）
                    expected_type = col_type.split(' ')[0].upper()
                      
                    if col_name not in existing_columns_dict:
                        self.add_column(table_name, col_name, col_type)
                    else:
                        # 获取实际类型
                        actual_type = existing_columns_dict[col_name].upper()
                         
                        # 对于VARCHAR类型，需要特殊处理以忽略长度差异（如果用户确认数据库类型正确）
                        if expected_type.startswith('VARCHAR') and actual_type.startswith('VARCHAR'):
                            # 都是VARCHAR类型，认为兼容，不提示类型不匹配
                            continue
                        # 对于其他类型，比较基本类型
                        elif expected_type.split('(')[0] != actual_type.split('(')[0]:
                            print(f"表 {table_name} 中的列 {col_name} 类型不匹配: 预期 {expected_type}, 实际 {actual_type}")
                            choice = input("是否修改该列类型？(y/n): ").lower()
                            if choice == 'y':
                                self.modify_column(table_name, col_name, col_type)
        
        self.disconnect()
        print("数据库生成完成")
        return True

if __name__ == '__main__':
    generator = MySQLDBGenerator()
    generator.generate()